package com.leal.util

import org.apache.log4j.Logger
import org.apache.spark.sql.Row

import java.sql.{Connection, DriverManager, PreparedStatement, SQLException}
import java.util

/**
 * @Author: leali
 * @Version: 1.0
 * @Date: 2021/9/8-17:56
 * @Description:
 */
object JdbcUtil {

  val logger: Logger = Logger.getLogger(this.getClass)

  def getConnection(driver: String, url: String, user: String, password: String): Connection = {
    Class.forName(driver)
    var connection: Connection = null
    try {
      connection = DriverManager.getConnection(url, user, password)
    } catch {
      case e: ClassNotFoundException => logger.error("Driver not loaded:", e)
      case e: SQLException => logger.error("Connection error: ", e)
      case e: Exception => e.printStackTrace()
    }
    logger.info(s"获得 连接$connection")
    connection
  }

  /**
   * after close
   *
   * @param connection 连接对象
   */
  def closeConnection(connection: Connection): Unit = {
    if (null != connection) {
      connection.close()
    }
  }

  /**
   * 重载
   *
   * @param mode NODE_MYSQL
   * @return
   */
  def getConnection(mode: String): Connection = {
    val envInfo: util.LinkedHashMap[String, String] = YamlUtil.getEnvInfo(mode)
    getConnection(driver = envInfo.get("driver"), url = envInfo.get("url"),
      user = envInfo.get("user"), password = envInfo.get("password"))
  }

  def operateStatement(connection: Connection, mode: String, sql: String, args: Any*): Any = {
    val statement: PreparedStatement = connection.prepareStatement(sql)
    var index = 1

    for (arg <- args) {
      arg match {
        case s: String => statement.setString(index, s)
        case i: Int => statement.setInt(index, i)
        case d: Double => statement.setDouble(index, d)
        case _ => println("args is error: ", arg)
      }
      index += 1
    }
    mode match {
      // executeQuery 用于查，可以返回ResultSet集
      case "SELECT" => statement.executeQuery()
      //executeUpdate 用于增删改，返回影响的行数
      case "INSERT" | "DELETE" | "UPDATE" => statement.executeUpdate()
      // TRUNCATE
      case _ => statement.execute()
    }
  }


  def upsertData(rows: Iterator[Row], mode: String, tableName: String, columns: Seq[String], batchSize: Int = 1000): Unit = {
    var connection: Connection = null
    var preparedStatement: PreparedStatement = null
    val sql = s"REPLACE INTO $tableName (${columns.mkString(", ")}) VALUES (${columns.map((_: String) => "?").mkString(", ")})"
    logger.info(s"sql $sql")
    try {
      connection = getConnection(mode)
      // get length sql
      /**
       * REPLACE 只会替换主键id 相同的纪录
       */
      preparedStatement = connection.prepareStatement(sql)
      connection.setAutoCommit(false)
      val iterator: rows.GroupedIterator[Row] = rows.sliding(batchSize, batchSize)
      while (iterator.hasNext) {
        val batch: Seq[Row] = iterator.next()
        batch.foreach { row: Row =>
          columns.zipWithIndex.foreach { case (column, index) => preparedStatement.setObject(index + 1, row.getAs(column)) }
          preparedStatement.addBatch()
        }
        logger.info(s"执行插入数据 batch size: ${batch.size}")
        preparedStatement.executeBatch()
      }
      connection.commit()
    } catch {
      case e: Exception => e.printStackTrace()
    } finally {
      preparedStatement.close()
      connection.close()
    }
  }
}
